import torch
from torch import nn

import d2l

# X = torch.tensor([
#     [0.0, 1.0, 2.0],
#     [3.0, 4.0, 5.0],
#     [6.0, 7.0, 8.0]
# ])
#
# K = torch.tensor([
#     [0.0, 1.0],
#     [2.0, 3.0]
# ])
#
# print(d2l.corr2d(X, K))


class Conv2D(nn.Module):
    def __init__(self, kernel_size):
        super().__init__()
        self.weight = nn.Parameter(torch.rand(kernel_size))
        self.bias = nn.Parameter(torch.zeros(1))

    def forward(self, X):
        return d2l.corr2d(X, self.weight) + self.bias


X = torch.ones((6, 8))
X[:, 2:6] = 0
K = torch.tensor([[1.0, -1.0]])

Y = d2l.corr2d(X, K)


conv2d = Conv2D((1, 2))

if True:
    conv2d = nn.Conv2d(1, 1, kernel_size=(1, 2), bias=False)
    X = X.reshape((1, 1, 6, 8))
    Y = Y.reshape((1, 1, 6, 7))

lr = 3e-2

for i in range(100):
    Y_hat = conv2d(X)
    print(Y_hat.shape, Y.shape)
    l = (Y_hat - Y) ** 2
    conv2d.zero_grad()
    l.sum().backward()
    conv2d.weight.data[:] -= lr * conv2d.weight.grad
    if (i + 1) % 2 == 0:
        print(f'epoch {i + 1}, loss {l.sum():.3f}')


print(type(conv2d))
print(conv2d.weight.data)
# print(conv2d.bias.data)
